#!/usr/bin/env python3

class Node:
    def __init__(self, left=None, key=None, right=None):
        self.height = 1
        self.left = left
        self.key = key
        self.right = right

    def preorder(self):
        print(str(self.key) + "(" + str(self.height) +  ")", end=' ')
        if self.left is not None:
            self.left.preorder()
        if self.right is not None:
            self.right.preorder()

    def inorder(self):
        if self.left is not None:
            self.left.inorder()
        print(str(self.key) + "(" + str(self.height) +  ")", end=' ')
        if self.right is not None:
            self.right.inorder()

    def print_graphviz(self, dot):
        if not self.left and not self.right:
            return
        code = str(hash(self))
        if self.left:
            dot.edge(str(self.key), str(self.left.key))
            self.left.print_graphviz(dot)
        else:
            name = code + 'left'
            dot.node(name, shape='point')
            dot.edge(str(self.key), name)

        if self.right:
            dot.edge(str(self.key), str(self.right.key))
            self.right.print_graphviz(dot)
        else:
            name = code + 'right'
            dot.node(name, shape='point')
            dot.edge(str(self.key), name)

    def recal_height(self):
        left_height = 0 if self.left is None else self.left.height
        right_height = 0 if self.right is None else self.right.height
        self.height = max(left_height, right_height) + 1
        


class AVLTree:
    def __init__(self):
        self.root = None


    def insert(self, key):
        node = Node(key=key)
        if self.root is None:
            self.root = node
        else:
            self.root = self._insert_internal(self.root, key)

    def _insert_internal(self, root, key):
        if root.key > key:
            # left
            if root.left is None:
                node = Node(key=key)
                root.left = node
            else:
                root.left = self._insert_internal(root.left, key)
        else:
            # right
            if root.right is None:
                node = Node(key=key)
                root.right = node
            else:
                root.right = self._insert_internal(root.right, key)
        left_height = 0 if root.left is None else root.left.height
        right_height = 0 if root.right is None else root.right.height
        root.height = max(left_height, right_height) + 1
        delta = left_height - right_height
        result_root = None
        if delta > 1:
            if root.left.key > key:
                # left left case
                result_root = self._right_rotate(root, root.left)
            else:
                # left right case
                root.left = self._left_rotate(root.left, root.left.right)
                result_root = self._right_rotate(root, root.left)
        elif delta < -1:
            if root.right.key > key:
                # right left case:
                root.right = self._right_rotate(root.right, root.right.left)
                result_root = self._left_rotate(root, root.right)
            else:
                # right right case
                result_root = self._left_rotate(root, root.right)
        else:
            result_root = root
        return result_root

    def delete(self, key):
        if self.root is None:
            return
        self.root = self._delete_internal(self.root, key)

    def _delete_internal(self, root, key):
        if root.key == key:
            min_node = self._find_min_node(root.right)
            if min_node:
                root.key = min_node.key
                root.right = self._delete_internal(root.right, min_node.key)
            else:
                # 叶节点
                return None
        elif root.key > key:
            root.left = self._delete_internal(root.left, key)
        else:
            root.right = self._delete_internal(root.right, key)

        root.recal_height()

        left_height = 0 if not root.left else root.left.height
        right_height = 0 if not root.right else root.right.height
        delta = left_height - right_height
        result_root = root
        if delta < -1:
            left_height = 0 if not root.right.left else root.right.left.height
            right_height = 0 if not root.right.right else root.right.right.height
            if root.right.height == right_height + 1:
                # right right
                result_root = self._left_rotate(root, root.right)
            else:
                # right left
                root.right = self._right_rotate(root.right, root.right.left)
                result_root = self._left_rotate(root, root.right)
        elif delta > 1:
            left_height = 0 if not root.left.left else root.left.left.height
            right_height = 0 if not root.left.right else root.left.right.height
            if root.left.height == left_height + 1:
                # left left
                result_root = self._right_rotate(root, root.left)
            else:
                # left right
                root.left = self._left_rotate(root.left, root.left.right)
                result_root = self._right_rotate(root, root.left)
        return result_root

        

    def _find_min_node(self, node):
        current = node
        while current and current.left:
            current = current.left
        return current

    
        

    def preorder(self):
        if self.root is None:
            print("AVL tree is empty")
        else:
            self.root.preorder()
        print("")

    def inorder(self):
        if self.root is None:
            print("AVL tree is empty")
        else:
            self.root.inorder()
        print("")

    def print_graphviz(self):
        if self.root is None:
            print("AVL tree is empty")
        else:
            from graphviz import Digraph
            dot = Digraph(comment='one AVL Tree')
            self.root.print_graphviz(dot)
            print(dot.source)
            dot.view()
            

    def _left_rotate(self, z, y):
        z.right = y.left
        y.left = z
        z.recal_height()
        y.recal_height()
        return y

    def _right_rotate(self, z, y):
        z.left = y.right
        y.right = z
        z.recal_height()
        y.recal_height()
        return y



# 测试插入
def test_insert_and_del():
    tree = AVLTree()
    for index in range(1, 8):
        tree.insert(index)
    tree.insert(9)
    tree.insert(15)
    tree.insert(16)

    print("after insert 1..10, 15, 16")
    tree.preorder()
    tree.inorder()

    print("\nafter delete 16:")
    tree.delete(16)
    tree.preorder()
    tree.inorder()


    print("\nafter delete 3:")
    tree.delete(3)
    tree.preorder()
    tree.inorder()

    tree.print_graphviz()


if __name__ == '__main__':
    test_insert_and_del()
